import datetime
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.structures import Instances, Boxes
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import torch
from detectron2.config import get_cfg
import cv2
import numpy as np
# from clear_directory import clear_directory


def setup_cfg_inference(model_dir, pre_config):
    """
    推理配置文件设置
    """
    cfg = get_cfg()
    cfg.merge_from_file(pre_config)  # 加载预定义的配置文件
    cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    cfg.MODEL.WEIGHT = model_dir + "/model_car.pth"  # 指定保存的模型文件路径
    cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # 根据你的数据集设置类别数量
    cfg.MODEL.ROI_HEADS.SCORE_THRESHOLD_TEST = 0.5  # 设置推理时的置信度阈值
    cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = 0.5  # 设置 Soft-NMS 的阈值
    cfg.DATASETS.TEST = ("ccpd_test",)  # 设置测试数据集
    return cfg


def load_model(cfg):
    """
    加载保存的模型
    """
    # 构建模型
    model = build_model(cfg)
    model.eval()  # 设置为推理模式

    # 加载模型权重
    # 使用 torch.load 并设置 weights_only=True
    with torch.no_grad():
        checkpoint = torch.load(cfg.MODEL.WEIGHT, map_location=cfg.MODEL.DEVICE, weights_only=True)
        model.load_state_dict(checkpoint["model"])
    return model


def inference_picture(model, picture_path):
    """
    对单张图像进行推理
    """
    # 加载图像
    image = cv2.imread(picture_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 进行推理
    with torch.no_grad():
        inputs = {"image": torch.as_tensor(image.astype("float32").transpose(2, 0, 1))}
        outputs = model([inputs])[0]

    instances = outputs["instances"].to("cpu")

    return instances


def picture_flag(model, picture_path, result_image_dir, show_image=False, only_max_score=True):
    """
    标记单张图像的车牌图片，可根据only_max_score选择是否只标记置信度最高的部分
    """
    image = cv2.imread(picture_path)
    instances = inference_picture(model, picture_path)

    # 如果没有检测到任何结果，直接返回
    if len(instances) == 0:
        print("未检测到任何车牌")
        return

    if only_max_score:
        # 获取置信度最高的检测结果的索引
        max_score_index = torch.argmax(instances.scores)
        # 提取置信度最高的检测结果的各个属性
        max_score = instances.scores[max_score_index]
        max_pred_boxes = instances.pred_boxes.tensor[max_score_index].unsqueeze(0)
        max_pred_masks = instances.pred_masks[max_score_index].unsqueeze(0)
        max_pred_classes = instances.pred_classes[max_score_index].unsqueeze(0)
        # 创建一个新的 Instances 对象，只包含置信度最高的检测结果
        instances = Instances(
            image_size=instances.image_size,
            scores=max_score.unsqueeze(0),
            pred_boxes=Boxes(max_pred_boxes),
            pred_masks=max_pred_masks,
            pred_classes=max_pred_classes
        )

    # 手动创建 Metadata 对象
    metadata = MetadataCatalog.get("ccpd_test")
    metadata.thing_classes = ["license_plate"]  # 设置类别名称
    # metadata.thing_colors = [(255, 0, 0)]  # 手动设置类别颜色（可选）

    # 可视化结果
    v = Visualizer(image[:, :, ::-1], metadata, scale=1.2)
    out = v.draw_instance_predictions(instances)
    result_image = out.get_image()[:, :, ::-1]  # BGR格式

    now_time = datetime.datetime.now().strftime('%m%d%H%M%S')
    result_image_name = result_image_dir + f"/picture_flag/{now_time}_0.jpg"
    # 保存或显示结果
    cv2.imwrite(result_image_name, result_image)
    print(f"标记车牌后的图像已保存到: {result_image_name}")
    if show_image:
        cv2.imshow("Result", result_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    # 返回RGB格式的结果图像
    return result_image[:, :, ::-1]


def picture_flag_license_plate(model, picture_path, result_image_dir):
    """
    仅提取保存车牌部分图片，并提取置信度最高的车牌部分
    """
    # 加载图像
    image = cv2.imread(picture_path)
    instances = inference_picture(model, picture_path)

    # 如果没有检测到任何结果，直接返回
    if len(instances) == 0:
        print("未检测到任何车牌")
        return

    # 获取置信度最高的检测结果
    max_score_index = torch.argmax(instances.scores)
    pred_mask = instances.pred_masks[max_score_index]
    bbox = instances.pred_boxes.tensor[max_score_index].numpy().astype(int)
    x1, y1, x2, y2 = bbox

    # 先用边界框裁剪图像
    cropped_image = image[y1:y2, x1:x2]

    # 保存裁剪后的图像
    now_time = datetime.datetime.now().strftime('%m%d%H%M%S')
    cropped_image_name = result_image_dir + f"/picture_flag_license_plate/{now_time}_1_1.jpg"
    cv2.imwrite(cropped_image_name, cropped_image)
    print(f"裁剪后的车牌图片已保存到: {cropped_image_name}")

    # 再用分割掩码提取车牌部分
    cropped_mask = pred_mask.numpy().astype(np.uint8) * 255
    cropped_mask = cropped_mask[y1:y2, x1:x2]
    license_plate = cropped_image * cropped_mask[:, :, np.newaxis]

    # 保存车牌部分的图像
    now_time = datetime.datetime.now().strftime('%m%d%H%M%S')
    license_plate_name = result_image_dir + f"/picture_flag_license_plate/{now_time}_1.jpg"
    cv2.imwrite(license_plate_name, license_plate[:, :, ::-1])
    print(f"置信度最高的识别的车牌部分已保存到: {license_plate_name}")

    return license_plate



if __name__ == "__main__":
    model_dir = "model_output"  # 模型保存目录
    pre_config = "configs/1/mask_rcnn_R_50_FPN_3x.yaml"  # 配置文件路径
    cfg = setup_cfg_inference(model_dir, pre_config)
    # 加载模型
    model = load_model(cfg)
    print("模型加载成功.")

    result_image_dir = "picture_inference"
    test_image_path = "test_img/img1.jpeg"

    # clear_directory(result_image_dir + "/picture_flag", auto_confirm=True)
    # clear_directory(result_image_dir + "/picture_flag_license_plate", auto_confirm=True)

    picture_flag(model, test_image_path, result_image_dir)
    picture_flag_license_plate(model, test_image_path, result_image_dir)
